import logging
import math
import os
import sys

import joblib
import numpy as np
import pandas as pd
from PyQt6.QtCore import Qt, QThread, pyqtSignal
from PyQt6.QtGui import QTextCursor, QPixmap
from PyQt6.QtWidgets import QMainWindow, QFileDialog, QApplication, QGraphicsPixmapItem, QGraphicsScene, QMessageBox
from matplotlib import pyplot as plt
from sklearn import metrics
from sklearn.model_selection import train_test_split
from xgboost import plot_importance

from App import Ui_MainWindow
from ReadExcel import GetCol_x_y
from ReadExcel import GetDevice_x_y
from ReadExcel import GetMode_x_y
from TrainModel import train_and_save_model, augment_data

globe_default_input = [20, 37, 20, 21, 10, 18, 28, 1, 6.8]
device_name = ['device_60_1', 'device_60_2', 'device_60_3', 'device_60_4',
               'device_50_1', 'device_50_2', 'device_50_3', 'device_50_4', 'device_50_5',
               'device_50_6', 'device_30_1', 'device_30_2', 'device_30_3', 'device_30_4']


class TrainThread(QThread):
    finished = pyqtSignal()
    progress = pyqtSignal(int)
    log_signal = pyqtSignal(str)

    def __init__(self, data_dict, train_input, augment_num, noise_num, testsize, model_path, param_grid,
                 cv_num, scoring):
        super().__init__()
        self.data_dict = data_dict
        self.train_input = train_input
        self.augment_num = augment_num
        self.noise_num = noise_num
        self.testsize = testsize
        self.model_path = model_path
        self.param_grid = param_grid
        self.cv_num = cv_num
        self.scoring = scoring

    def run(self):
        self.__slip_and_train_model()

    def Evaluate(self, y_test, y_pred):
        # 计算评估指标
        explained_variance = metrics.explained_variance_score(y_test, y_pred)
        mean_absolute_error = metrics.mean_absolute_error(y_test, y_pred)
        mean_squared_error = metrics.mean_squared_error(y_test, y_pred)
        r2_score = metrics.r2_score(y_test, y_pred)
        # 记录评估结果到日志
        self.log_signal.emit('  可解释方差：{}'.format(round(explained_variance, 10)))
        self.log_signal.emit('  平均绝对误差：{}'.format(round(mean_absolute_error, 10)))
        self.log_signal.emit('  均方误差：{}'.format(round(mean_squared_error, 10)))
        self.log_signal.emit('  R方差：{}'.format(round(r2_score, 10)))

        num_outputs = y_test.shape[1]
        fig, axes = plt.subplots(num_outputs, 1, figsize=(10, 2 * num_outputs), sharex=True)

        for i in range(num_outputs):
            ax = axes[i] if num_outputs > 1 else axes
            ax.plot(y_test[:, i], label='Actual', marker='o')
            ax.plot(y_pred[:, i], label='Predicted', marker='x')
            ax.set_ylabel(f'Output {i + 1}')
            ax.legend()
            ax.grid(True)

        plt.xlim(0, 50)  # 仅查看前50个样本
        plt.xlabel('Sample Index')
        plt.suptitle('Model Predictions vs Actual Values')
        plt.tight_layout(rect=[0, 0, 1, 0.96])

        output_dir = './Visualization'
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        # 保存图像到指定路径
        plt.savefig(os.path.join(output_dir, 'train_plt.png'))
        plt.close()

    def __slip_and_train_model(self):
        split_data = {}
        total_keys = len(self.data_dict)

        # 初始化整体图形用于显示所有模型的重要性图
        combined_fig, combined_axes = plt.subplots(total_keys, 1, figsize=(8, 4 * total_keys))

        if total_keys == 1:
            combined_axes = [combined_axes]  # 确保axes总是列表形式，便于迭代

        for idx, (key, y) in enumerate(self.data_dict.items()):
            x_aug, y_aug = augment_data(np.array(self.train_input), y, self.augment_num, self.noise_num)
            self.log_signal.emit(f'{key}模型数据增强成功，正在训练，请稍等。')
            x_train, x_test, y_train, y_test = train_test_split(x_aug, y_aug, test_size=self.testsize, random_state=42)
            split_data[f'x_train_{key}'] = x_train
            split_data[f'y_train_{key}'] = y_train
            split_data[f'x_test_{key}'] = x_test
            split_data[f'y_test_{key}'] = y_test

            y_pred, best_model, best_param = train_and_save_model(x_train, y_train, x_test, y_test, key,
                                                                  self.model_path,
                                                                  self.param_grid, self.cv_num, self.scoring)
            self.log_signal.emit(f'{key}模型的最佳参数为{best_param}')

            best_xgb_model = best_model.named_steps['regressor'].estimators_[0]

            # 在子图上绘制特征重要性图
            ax = combined_axes[idx]
            plot_importance(best_xgb_model, importance_type='weight', ax=ax)
            ax.set_title(f'{key} Feature Importance')

            self.Evaluate(y_test, y_pred)
            self.progress.emit(int((idx + 1) / total_keys * 100))

        plt.tight_layout()
        plt.savefig(r'./Visualization/combined_importance.png')
        plt.close(combined_fig)

        self.finished.emit()


class PredictThread(QThread):
    progress_signal = pyqtSignal(int)
    load_models_signal = pyqtSignal(str)  # 用于传递模型路径等信息
    frequency_signal = pyqtSignal(int)
    result_signal = pyqtSignal(np.ndarray)
    finished_signal = pyqtSignal()  # 用于通知线程完成

    def __init__(self, ui):
        super().__init__()
        self.ui = ui

    def run(self):
        current_text = self.ui.use_fun_combo.currentText()
        input_list = [
            round(self.ui.input_100.value(), 2),
            round(self.ui.input_100_200.value(), 2),
            round(self.ui.input_200_325.value(), 2),
            round(self.ui.input_325.value(), 2),
            round(self.ui.input_Dv10.value(), 2),
            round(self.ui.input_Dv50.value(), 2),
            round(self.ui.input_Dv90.value(), 2),
            round(self.ui.input_zhenshi.value(), 3),
            round(self.ui.input_bibiao.value(), 4)
        ]

        model_path = {
            "按设备": './Model/ByDevice',
            "按模块": './Model/ByMode',
            "按机器": './Model/ByCol'
        }.get(current_text, './Model')
        devices_name_dict = set(device_name)
        devices_name_dict.add('feed_frequency')

        def load_models(keys):
            return {key: joblib.load(os.path.join(model_path, f'{key}_model.pkl')) for key in keys}

        model_keys = {
            "按设备": [f'device_{n}_{i}' for n in ['60', '50', '30'] for i in range(1, 7) if
                       f'device_{n}_{i}' in devices_name_dict],
            "按模块": [f'{mode}_{n}' for mode in ['main', 'neifen', 'waifen', 'fengji'] for n in ['60', '50', '30']],
            "按机器": ['main', 'neifen', 'waifen', 'fengji']
        }[current_text]

        try:
            model_dict = load_models(model_keys)
        except FileNotFoundError as e:
            self.load_models_signal.emit(f'未找到模型，请先进行训练。Error:{e}')
            return  # 退出方法以避免后续错误
        finally:
            self.finished_signal.emit()  # 发送完成信号

        input_array = np.array(input_list).reshape(1, -1)
        predictions_array = []

        num_models = len(model_dict)  # 预先获取模型数量
        for i, (module, model) in enumerate(model_dict.items()):
            y_pred_param = model.predict(input_array)
            predictions_array.append(np.clip(np.round(y_pred_param), 0, 99).astype(int)[0].tolist())

            # 更新进度信号
            self.progress_signal.emit(int((i + 1) / num_models * 100))

        pred_fun = self.ui.use_fun_combo.currentText()

        if pred_fun == "按设备":
            pred = np.array(predictions_array)
            feed_model = joblib.load('./Model/ByDevice/feed_frequency_model.pkl')
            feed_frequency = feed_model.predict(input_array)
        elif pred_fun == "按模块":
            pred = np.array([sum(predictions_array[i:i + 3], []) for i in range(0, len(predictions_array), 3)]).T
            feed_model = joblib.load('./Model/ByMode/feed_frequency_model.pkl')
            feed_frequency = feed_model.predict(input_array)
        else:
            pred = np.array(predictions_array).T
            feed_model = joblib.load('./Model/ByCol/feed_frequency_model.pkl')
            feed_frequency = feed_model.predict(input_array)
        int_feed_frequency = np.clip(np.round(feed_frequency), 0, 99).astype(int).item()

        self.result_signal.emit(pred)
        self.frequency_signal.emit(int_feed_frequency)


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)

        self.ui.SelectPath_Button.clicked.connect(self.__select_button_click)
        self.ui.Read_Button.clicked.connect(self.__read_button_click)

        self.ui.refresh_brower_btn.clicked.connect(
            lambda: self.__add_image_to_scene('./Visualization/train_plt.png', 'vis_wid'))
        self.ui.refresh_brower_btn_2.clicked.connect(
            lambda: self.__add_image_to_scene('./Visualization/combined_importance.png', 'importance_vis'))
        self.ui.Train_Button.clicked.connect(self.__train_button_click)
        self.ui.Use_Button.clicked.connect(self.__predict_parameters)
        self.ui.Export_Button.clicked.connect(self.__export_data)
        self.train_log_file = 'TrainLog.log'
        self.pred_log_file = 'PredictLog.log'
        self.__setup_logging()  # 只在初始化时设置日志
        self.train_thread = None
        self.Pred_thread = None

    def __select_button_click(self):
        folder_path = QFileDialog.getExistingDirectory(self, "选择文件夹")
        if folder_path:
            self.ui.DataPath.setText(folder_path)

    def __setup_logging(self):
        # 清除已有的全局处理器
        for handler in logging.root.handlers[:]:
            logging.root.removeHandler(handler)

        # 禁用默认的控制台日志
        logging.getLogger().setLevel(logging.WARNING)
        logging.getLogger().addHandler(logging.NullHandler())  # 阻止控制台输出

        # 训练日志配置
        self.train_logger = logging.getLogger('train_logger')
        self.train_logger.setLevel(logging.INFO)
        train_handler = logging.FileHandler(self.train_log_file)
        train_handler.setLevel(logging.INFO)
        train_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        train_handler.setFormatter(train_formatter)
        self.train_logger.addHandler(train_handler)

        # 预测日志配置
        self.pred_logger = logging.getLogger('pred_logger')
        self.pred_logger.setLevel(logging.INFO)
        pred_handler = logging.FileHandler(self.pred_log_file)
        pred_handler.setLevel(logging.INFO)
        pred_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        pred_handler.setFormatter(pred_formatter)
        self.pred_logger.addHandler(pred_handler)

    @staticmethod
    def __update_log(log_file, text_area):
        with open(log_file, "r") as file:
            log_content = file.read()
        text_area.setPlainText(log_content)

        cursor = text_area.textCursor()
        cursor.movePosition(QTextCursor.MoveOperation.End)
        text_area.setTextCursor(cursor)
        text_area.ensureCursorVisible()

    def __update_train_log(self):
        self.__update_log(self.train_log_file, self.ui.train_info)

    def __update_pred_log(self):
        self.__update_log(self.pred_log_file, self.ui.use_info)

    def log_train_info(self, message):
        self.train_logger.info(message)
        self.__update_train_log()  # 更新训练日志显示

    def log_pred_info(self, message):
        self.pred_logger.info(message)
        self.__update_pred_log()  # 更新预测日志显示

    @staticmethod
    def __prepare_data_dict(read_way, *data):
        keys = {
            "按设备": ['feed_frequency'] + device_name,
            "按模块": ['feed_frequency', 'main_60', 'main_50', 'main_30', 'neifen_60', 'neifen_50', 'neifen_30',
                       'waifen_60', 'waifen_50', 'waifen_30', 'fengji_60', 'fengji_50', 'fengji_30'],
            "按机器": ['feed_frequency',
                       'main', 'neifen', 'waifen', 'fengji']
        }
        return dict(zip(keys[read_way], list(data)))

    def __read_button_click(self):
        data_path = self.ui.DataPath.toPlainText()
        read_way = self.ui.model_combo.currentText()

        get_data_func = {
            "按设备": GetDevice_x_y,
            "按模块": GetMode_x_y,
            "按机器": GetCol_x_y
        }.get(read_way)

        try:
            train_input, *data, log_message = get_data_func(data_path, globe_default_input)
            data_dict = self.__prepare_data_dict(read_way, *data)
            self.log_train_info(f'数据空缺：{log_message}')
            self.log_train_info(f'成品输入：{train_input}')
            for key, value in data_dict.items():
                self.log_train_info(f"{key}: {value}")
        except Exception as e:
            self.log_train_info(f"在读取或处理数据时发生错误: {e}")

    def __train_button_click(self):
        if self.train_thread is not None:
            self.log_train_info('正在训练其他模型，请稍等')
            return

        try:
            current_text = self.ui.model_combo.currentText()
            model_path = self.__get_model_paths(current_text)
            get_data_func = self.__get_data_function(current_text)
            data_path = self.ui.DataPath.toPlainText()

            # Extract parameters and prepare data
            train_input, *data, logmess = get_data_func(data_path, globe_default_input)
            data_dict = self.__prepare_data_dict(current_text, *data)

            # Prepare and start thread
            self.__prepare_and_start_thread(train_input, data_dict, model_path)
        except ValueError as ve:
            self.log_train_info(f"Validation error: {ve}")
        except Exception as e:
            self.log_train_info(f"An error occurred: {e}")

    @staticmethod
    def __get_model_paths(current_text):
        model_paths = {
            "按设备": r'./Model/ByDevice',
            "按模块": r'./Model/ByMode',
            "按机器": r'./Model/ByCol'
        }
        return model_paths.get(current_text, (None, None))

    @staticmethod
    def __get_data_function(current_text):
        get_data_func = {
            "按设备": GetDevice_x_y,
            "按模块": GetMode_x_y,
            "按机器": GetCol_x_y
        }
        return get_data_func.get(current_text)

    def __get_scoring_method(self):
        score_way = self.ui.ScoreCombo.currentText()
        get_score_way = {
            "MSE": 'neg_mean_squared_error',
            "R2": 'r2',
            "MAE": 'neg_mean_absolute_error'
        }
        return get_score_way.get(score_way)

    def __prepare_and_start_thread(self, train_input, data, model_path):
        param_grid = self.__build_param_grid()
        augment_num = self.ui.augment_num.value()
        noise_num = self.ui.noise_spin.value()
        testsize = self.ui.testsize_dspin.value()
        cv_num = self.ui.CVSpin.value()
        scoring = self.__get_scoring_method()

        # 创建训练线程实例
        self.train_thread = TrainThread(data, train_input, augment_num, noise_num, testsize, model_path,
                                        param_grid, cv_num, scoring)
        # 连接训练线程的日志信号到更新日志的方法，以便在训练过程中接收日志信息
        self.train_thread.log_signal.connect(self.update_thread_log)
        # 连接训练线程的进度信号到更新进度的方法，以便在训练过程中接收进度更新
        self.train_thread.progress.connect(self.update_progress)
        # 连接训练线程完成信号到处理训练完成的方法，以便在训练结束后更新UI或执行其他操作
        self.train_thread.finished.connect(self.training_complete)
        # 连接完成信号到删除线程的方法，确保线程在完成后被安全地删除，以释放资源
        self.train_thread.finished.connect(self.train_thread.deleteLater)
        # 启动训练线程，开始模型训练过程
        self.train_thread.start()

    def update_thread_log(self, message):
        # 确保这个方法在主线程中执行
        self.log_train_info(message)

    def training_complete(self):
        self.update_progress(100)  # 确保进度条满
        QMessageBox.information(self, "训练完成", "模型训练已经完成！")
        self.train_thread = None  # 重置为 None，允许重新启动预测

    def update_progress(self, value):
        """更新进度条"""
        if 0 <= value <= 100:
            self.ui.progressBar.setValue(value)
        else:
            self.log_train_info(f"Invalid progress value: {value}")

    def __build_param_grid(self):
        def convert_text_to_list(text, convert_func, default=None):
            try:
                return [convert_func(x.strip()) for x in text.split(',') if x.strip()] or default
            except ValueError as e:
                self.log_train_info(f"Error in __build_param_grid: {e}.  Use Default Param_grid：{default}")
                return default

        return {
            'regressor__estimator__n_estimators': convert_text_to_list(self.ui.n_estimators_edit.text(), int, [100]),
            'regressor__estimator__max_depth': convert_text_to_list(self.ui.MaxDepth_edit.text(), int, [3]),
            'regressor__estimator__learning_rate': convert_text_to_list(self.ui.LearningRate_Edit.text(), float, [0.1]),
            'regressor__estimator__alpha': [0.1],
            'regressor__estimator__lambda': convert_text_to_list(self.ui.L2_edit.text(), int, [1])
        }

    def __add_image_to_scene(self, image_path, wid_name):
        widget = getattr(self.ui, wid_name)
        # 创建一个 QGraphicsScene 对象
        self.scene = QGraphicsScene(self)
        widget.setScene(self.scene)
        self.scene.clear()
        self.log_train_info('可视化已刷新')
        pixmap = QPixmap(image_path)
        if not pixmap.isNull():
            item = QGraphicsPixmapItem(pixmap)
            item.setTransformationMode(Qt.TransformationMode.SmoothTransformation)
            self.scene.addItem(item)
            widget.setSceneRect(item.boundingRect())
            widget.fitInView(item, Qt.AspectRatioMode.KeepAspectRatio)
        else:
            self.log_train_info(f"Failed to load image: {image_path}")

    def wheelEvent(self, event):
        if event.modifiers() & Qt.KeyboardModifier.ControlModifier:
            scaleFactor = math.pow(1.1, event.angleDelta().y() / 120)
            self.scaleView(scaleFactor)
            event.accept()
        else:
            super().wheelEvent(event)

    def scaleView(self, scaleFactor):
        factor = self.ui.vis_wid.transform().m11() * scaleFactor
        if 0.07 < factor < 100:
            self.ui.vis_wid.scale(scaleFactor, scaleFactor)
        factor = self.ui.importance_vis.transform().m11() * scaleFactor
        if 0.07 < factor < 100:
            self.ui.importance_vis.scale(scaleFactor, scaleFactor)

    def __predict_parameters(self):
        if self.train_thread is not None:
            self.log_pred_info("正在训练模型，请等待")
            QMessageBox.warning(self, "警告", "正在训练模型，请等待", QMessageBox.StandardButton.Ok)
            return
        if self.Pred_thread is not None:
            self.log_pred_info("正在预测参数，请等待")
            QMessageBox.warning(self, "警告", "正在预测模型，请等待", QMessageBox.StandardButton.Ok)
            return

        self.Pred_thread = PredictThread(self.ui)
        self.Pred_thread.load_models_signal.connect(self.log_pred_info)
        self.Pred_thread.progress_signal.connect(self.ui.use_progressBar.setValue)
        self.Pred_thread.result_signal.connect(self.__print_parameters)
        self.Pred_thread.frequency_signal.connect(self.__print_feed_frequency)
        # 连接完成信号
        self.Pred_thread.finished_signal.connect(self.on_prediction_finished)

        # 连接线程结束信号，确保清理线程
        self.Pred_thread.finished.connect(self.Pred_thread.deleteLater)
        self.Pred_thread.start()

    def on_prediction_finished(self):
        # 清理或重置线程资源
        self.Pred_thread = None  # 重置为 None，允许重新启动预测

    def __print_parameters(self, prediction):
        self.log_pred_info(f'预测结果为{prediction}')
        for row in range(14):
            for col in range(4):
                widget = getattr(self.ui, f'col_{col + 1}_row_{row + 1}')
                widget.setText(str(prediction[row][col]))

    def __print_feed_frequency(self, frequency):
        self.log_pred_info(f'进料频率结果为{frequency}')
        self.ui.feed_frequency_result.setText(str(frequency))

    def __export_data(self):
        if self.train_thread is not None:
            self.log_pred_info("正在训练模型，请等待")
            return
        if self.Pred_thread is not None:
            self.log_pred_info("正在预测参数，请等待")
            return
        export_list = []  # 初始化为空列表
        first_cell_widget = getattr(self.ui, 'col_1_row_1')
        first_cell_text = first_cell_widget.text()  # 获取其文本内容

        # 检查 col_1_row_1 是否为空
        if not first_cell_text:
            self.log_pred_info("请先选择模型进行预测")
            QMessageBox.warning(self, "警告", "请先选择模型进行预测", QMessageBox.StandardButton.Ok)
            return  # 结束函数

        for row in range(14):
            current_row = []  # 创建当前行的列表
            for col in range(4):
                widget = getattr(self.ui, f'col_{col + 1}_row_{row + 1}')
                text = widget.text()
                current_row.append(text)
            export_list.append(current_row)  # 将当前行添加到 export_list

        try:
            prediction_df = pd.DataFrame(export_list, columns=['主机', '内分', '外分', '风机'])
            prediction_df.insert(0, '设备型号', pd.Series(device_name))
            file_path, _ = QFileDialog.getSaveFileName(self, "保存文件", "", "Excel 文件 (*.xlsx);;所有文件 (*)")
            if file_path:
                prediction_df.to_excel(file_path, index=False)
                self.log_pred_info(f"文件已保存到: {file_path}")
            else:
                self.log_pred_info("文件导出取消。")
        except Exception as e:
            self.log_pred_info(f"导出数据时出现错误: {e}")

    def closeEvent(self, event):
        if self.train_thread is not None or self.Pred_thread is not None:
            reply = QMessageBox.question(self, '确认', '有正在进行的任务。确定要退出吗？',
                                         QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
                                         QMessageBox.StandardButton.No)
            if reply == QMessageBox.StandardButton.Yes:
                if self.train_thread is not None:
                    self.train_thread.terminate()  # 安全地停止训练线程
                    self.train_thread.wait()  # 等待线程完全停止
                if self.Pred_thread is not None:
                    self.Pred_thread.terminate()  # 安全地停止预测线程
                    self.Pred_thread.wait()  # 等待线程完全停止
                event.accept()  # 现在可以安全关闭
            else:
                event.ignore()  # 不关闭窗口
        else:
            event.accept()  # 没有正在运行的任务，可以安全关闭


if __name__ == '__main__':
    app = QApplication(sys.argv)
    mainWindow = MainWindow()
    mainWindow.show()
    sys.exit(app.exec())
